/*
 * Copyright 2011 Atteo.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */
package org.atteo.classindex;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;

import org.apache.maven.plugins.shade.resource.ResourceTransformer;
import org.codehaus.plexus.util.IOUtil;

public class ClassIndexTransformer implements ResourceTransformer {
	public static final String SUBCLASS_INDEX_PREFIX = "META-INF/services/";
	public static final String ANNOTATED_INDEX_PREFIX = "META-INF/annotations/";
	public static final String PACKAGE_INDEX_NAME = "jaxb.index";

	private ByteArrayOutputStream data;

	private Map serviceEntries = new HashMap();

	@Override
	public boolean canTransformResource(String resource) {
		if (resource.startsWith(SUBCLASS_INDEX_PREFIX) || resource.startsWith(ANNOTATED_INDEX_PREFIX)
				|| resource.endsWith("/" + PACKAGE_INDEX_NAME)) {
			data = (ByteArrayOutputStream) serviceEntries.get(resource);

			if (data == null) {
				data = new ByteArrayOutputStream();
				serviceEntries.put(resource, data);
			}
			return true;
		}

		return false;
	}

	@Override
	public void processResource(String resource, InputStream is, List relocators)
			throws IOException {
		IOUtil.copy(is, data);
		is.close();
	}

	@Override
	public boolean hasTransformedResource() {
		return serviceEntries.size() > 0;
	}

	@Override
	public void modifyOutputStream(JarOutputStream jos)
			throws IOException {
		for (Iterator i = serviceEntries.keySet().iterator(); i.hasNext();) {
			String key = (String) i.next();
			ByteArrayOutputStream stream = (ByteArrayOutputStream) serviceEntries.get(key);
			jos.putNextEntry(new JarEntry(key));
			IOUtil.copy(new ByteArrayInputStream(stream.toByteArray()), jos);
			stream.reset();
		}
	}
}
